import torch
import torchvision
from torchvision.ops import nms  # 无报错即表示正常

print(torch.__version__)         # 输出 PyTorch 版本（如 1.12.0）
print(torchvision.__version__)   # 输出 TorchVision 版本（如 0.13.0）


# 验证 PyTorch 的 GPU 访问
print(f"CUDA 可用: {torch.cuda.is_available()}")  # 必须返回 True
print(f"当前 GPU 设备: {torch.cuda.get_device_name(0)}")  # 应显示 RTX 4060

# 创建一个 GPU 张量测试
test_tensor = torch.randn(1000, 1000).cuda()
print(f"张量所在设备: {test_tensor.device}")  # 应显示 cuda:0

# 执行一个 GPU 计算（观察 nvidia-smi 显存变化）
_ = test_tensor @ test_tensor.T
print(f"GPU 显存占用: {torch.cuda.memory_allocated(0)/1024**2:.2f} MB")